package com.yihaodian.search.grouping;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.lucene.document.Document;
import org.apache.lucene.index.IndexableField;
import org.apache.lucene.search.CachingCollector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.grouping.GroupDocs;
import org.apache.lucene.search.grouping.SearchGroup;
import org.apache.lucene.search.grouping.TopGroups;
import org.apache.lucene.search.grouping.term.TermAllGroupsCollector;
import org.apache.lucene.search.grouping.term.TermFirstPassGroupingCollector;
import org.apache.lucene.search.grouping.term.TermSecondPassGroupingCollector;
import org.apache.lucene.util.BytesRef;

import com.yihaodian.api.search.SearchException;
import com.yihaodian.api.search.SearchRequest;
import com.yihaodian.api.search.SearchResponse;
import com.yihaodian.api.search.grouping.GroupingResult;
import com.yihaodian.api.search.grouping.GroupingSpec;
import com.yihaodian.api.search.misc.SearchHit;
import com.yihaodian.search.SearchContext;
import com.yihaodian.search.highlight.HighlightHandler;

public class GroupingHandler {
	/**
	 * TODO: Sort groups by hits.
	 * @param context
	 * @throws SearchException
	 */
	public static void grouping(SearchContext context) throws SearchException {
		try {
			IndexSearcher searcher = context.getSearcher().searcher();
			SearchRequest searchRequest = context.getShardRequest().getSearchRequest();
			GroupingSpec spec = searchRequest.getGroupingSpec();
			if(spec == null)
				return;
			
			String groupingTerm = spec.getGroupingTerm();
			Sort groupSort = Sort.RELEVANCE;
			Sort docSort = Sort.RELEVANCE;

			Query query = context.getQuery();

			TermFirstPassGroupingCollector c1 = new TermFirstPassGroupingCollector(
					groupingTerm, groupSort, 10);

			boolean cacheScores = true;
			double maxCacheRAMMB = 4.0;
			CachingCollector cachedCollector = CachingCollector.create(c1,
					cacheScores, maxCacheRAMMB);
			searcher.search(query, context.getFilter(), cachedCollector);

			Collection<SearchGroup<BytesRef>> topGroups = c1.getTopGroups(0,
					true);

			if (topGroups == null) {
				// No groups matched
				return;
			}

			boolean getScores = true;
			boolean getMaxScores = true;
			boolean fillFields = true;
			TermSecondPassGroupingCollector c2 = new TermSecondPassGroupingCollector(
					groupingTerm, topGroups, groupSort, docSort, 10, getScores,
					getMaxScores, fillFields);

			// Optionally compute total group count
			TermAllGroupsCollector allGroupsCollector = null;
			if (true) {
				allGroupsCollector = new TermAllGroupsCollector(groupingTerm);
				// c2 = MultiCollector.wrap(c2, allGroupsCollector);
			}

			if (cachedCollector.isCached()) {
				// Cache fit within maxCacheRAMMB, so we can replay it:
				cachedCollector.replay(c2);
			} else {
				// Cache was too large; must re-execute query:
				searcher.search(query, context.getFilter(), c2);
			}

			TopGroups<BytesRef> groupsResult = c2.getTopGroups(0);
			if (true) {
				groupsResult = new TopGroups<BytesRef>(groupsResult,
						allGroupsCollector.getGroupCount());
			}

			// Generate groupsResult...
			GroupDocs<BytesRef>[] groups = groupsResult.groups;
			List<GroupingResult> results = new ArrayList<GroupingResult>();
			for (int i = 0; i < groups.length; i++) {
				GroupingResult result = new GroupingResult();
				GroupDocs<BytesRef> group = groups[i];
				if (group.groupValue == null) {
					continue;
				}
				result.setGroupName(group.groupValue.utf8ToString());
				ScoreDoc[] docs = group.scoreDocs;
				SearchHit[] hits = new SearchHit[docs.length];
				for (int j = 0; j < docs.length; j++) {
					SearchHit hit = new SearchHit();
					Document doc = searcher.getIndexReader().document(
							docs[j].doc);
					hit.setDocId(docs[j].doc);
					hit.setScore(docs[j].score);
					hits[j] = hit;
					Set<String> fields = searchRequest.getFields();
					if (fields != null && fields.size() > 0) {
						Map<String, String> fieldValues = new HashMap<String, String>();
						for (String field : fields) {
							fieldValues.put(field, doc.get(field));
						}
						hit.setFieldValues(fieldValues);
					} else {
						Map<String, String> fieldValues = new HashMap<String, String>();
						Iterator<IndexableField> it = doc.iterator();
						while (it.hasNext()) {
							IndexableField field = it.next();
							fieldValues.put(field.name(), field.stringValue());
						}
						hit.setFieldValues(fieldValues);
					}
				}
				result.setHits(hits);
				results.add(result);
			}
			
			Collections.sort(results);

			SearchResponse response = context.getResponse();
			if (response == null) {
				response = new SearchResponse();
			}
			response.setGroupingResult(results.toArray(new GroupingResult[0]));
		} catch (Exception e) {
			throw new SearchException(e);
		}
		
		postProcess(context);
	}
	
	public static void postProcess(SearchContext context) {
		HighlightHandler.highlightGrouping(context);
	}
}
